package com.aegis.common.web.xss;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.databind.node.TextNode;
import org.jsoup.Jsoup;
import org.jsoup.safety.Safelist;
import org.springframework.util.StreamUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;

/**
 *
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private final Safelist safelist;
    private final byte[] cachedBody;

    public XssHttpServletRequestWrapper(HttpServletRequest request, Safelist safelist) throws IOException {
        super(request);
        this.safelist = safelist;

        // 判断是否是 JSON 请求
        if (request.getContentType() != null && request.getContentType().toLowerCase().contains("application/json")) {
            String body = StreamUtils.copyToString(request.getInputStream(), StandardCharsets.UTF_8);
            String sanitizedBody = sanitizeJson(body);
            this.cachedBody = sanitizedBody.getBytes(StandardCharsets.UTF_8);
        } else {
            this.cachedBody = null;
        }
    }

    /**
     * 处理json数据
     * @param json
     * @return
     */
    private String sanitizeJson(String json) {
        try {
            ObjectMapper mapper = new ObjectMapper();
            JsonNode root = mapper.readTree(json);
            sanitizeNode(root);
            return mapper.writeValueAsString(root);
        } catch (Exception e) {
            return json; // 出现解析异常时返回原始 JSON
        }
    }

    private void sanitizeNode(JsonNode node) {
        if (node.isObject()) {
            ObjectNode objectNode = (ObjectNode) node;
            Iterator<Map.Entry<String, JsonNode>> fields = objectNode.fields();
            while (fields.hasNext()) {
                Map.Entry<String, JsonNode> entry = fields.next();
                JsonNode child = entry.getValue();
                if (child.isTextual()) {
                    String clean = Jsoup.clean(child.asText(), safelist);
                    objectNode.put(entry.getKey(), clean);
                } else {
                    sanitizeNode(child); // 递归清洗
                }
            }
        } else if (node.isArray()) {
            ArrayNode arrayNode = (ArrayNode) node;
            for (int i = 0; i < arrayNode.size(); i++) {
                JsonNode child = arrayNode.get(i);
                if (child.isTextual()) {
                    String clean = Jsoup.clean(child.asText(), safelist);
                    arrayNode.set(i, new TextNode(clean));
                } else {
                    sanitizeNode(child); // 递归
                }
            }
        }
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (cachedBody == null) {
            return super.getInputStream();
        }
        ByteArrayInputStream bais = new ByteArrayInputStream(cachedBody);
        return new ServletInputStream() {
            @Override public boolean isFinished() { return bais.available() == 0; }
            @Override public boolean isReady() { return true; }
            @Override public void setReadListener(ReadListener readListener) {}
            @Override public int read() { return bais.read(); }
        };
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    /**
     * 表单参数才会触发
     * @param name
     * @return
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values != null) {
            int length = values.length;
            String[] escapseValues = new String[length];
            for (int i = 0; i < length; i++) {
                escapseValues[i] = Jsoup.clean(values[i], safelist).trim();
            }
            return escapseValues;
        }
        return super.getParameterValues(name);
    }
}